'''
Extends TelegramHandler by managing the database and adding support for
multiple sessions
'''

from operator import itemgetter
from hashlib import sha256
from os import path
from sys import platform

from backend import utils
from backend.telegramHandler import TelegramHandler
from backend.sHandler_extern import SHandler_extern

class SessionsHandler(SHandler_extern):
    def __init__(self, notif_fun: callable, local_library: bool = False):
        super().__init__()

        self.tgSessions = {}
        self.freeSessions = []
        self.notif_fun = notif_fun

        self.asyncFiles = utils.AsyncFiles(
            "{}transferHandler_extern.{}".format('../' if local_library else '',
            'dll' if platform == 'win32' else 'so'))

        for sFile in range(1, int(self.cfg['remote']['max_sessions'])+1):
            # set session as free only if there is no resume info for it
            if not self.resumeData[sFile]:
                self.freeSessions.append(sFile)

            self.tgSessions[sFile] = {}

            # initialize all sessions that will be used
            self.tgSessions[sFile]['fun'] = TelegramHandler(
                self.cfg, sFile, self._saveProgress,
                self.saveResumeData, self.asyncFiles)

            self.tgSessions[sFile]['info'] = {}
            self.tgSessions[sFile]['info']['rPath'] = ''
            self.tgSessions[sFile]['info']['progress'] = 0
            self.tgSessions[sFile]['info']['size'] = 0
            self.tgSessions[sFile]['info']['type'] = None


    '''
    Returns:
    The first available session
    The given one
    -1 if there's no free session
    '''
    def _useSession(self, sFile: int = None):
        if sFile:
            '''
            Don't remove session because it was already removed
            in resumeHandler
            '''
            return sFile

        if not self.freeSessions:
            self.notif_fun("All sessions are currently used")
            return -1

        # get available session
        retSession = self.freeSessions[0]
        self.freeSessions.pop(0)
        return retSession


    def _freeSession(self, sFile: int):
        if not sFile in range(1, int(self.cfg['remote']['max_sessions']) + 1):
            raise IndexError("sFile should be between 1 and {}.".format(
                int(self.cfg['remote']['max_sessions'])
            ))

        if sFile in self.freeSessions:
            raise ValueError("Can't free a session that is already free.")

        self.freeSessions.append(sFile)


    def _saveProgress(self, current, total, current_chunk, total_chunks, sFile):
        prg = int(((current/total/total_chunks)+
                   (current_chunk/total_chunks))*100)

        self.tgSessions[sFile]['info']['progress'] = prg


    async def resumeHandler(self, sFile: int, selected: int):
        if not sFile in range(1, int(self.cfg['remote']['max_sessions']) + 1):
            raise IndexError("sFile should be between 1 and {}.".format(
                int(self.cfg['remote']['max_sessions'])
            ))

        if selected == 1: # finish the transfer
            if self.resumeData[sFile]['type'] == 'upload':
                await self.upload(self.resumeData[sFile], sFile)
            elif self.resumeData[sFile]['type'] == 'download':
                await self.download(self.resumeData[sFile], sFile)

        elif selected == 2: # delete the resume file
            if self.resumeData[sFile]['type'] == 'upload':
                await self._cleanTg(self.resumeData[sFile]['fileID'], sFile)
            else:
                self._freeSession(sFile)

            self.delResumeData(sFile)


    '''
    Returns:
    -1 if no free sessions
    0 if successfull
    '''
    async def _cleanTg(self, IDList: list = [], sFile: int = None) -> int:
        sFile = self._useSession(sFile)
        if sFile == -1: return -1

        mode = 2

        if not IDList:
            mode = 1
            for i in self.fileDatabase['file']:
                for j in i['fileID']:
                    IDList.append(j)

        await self.tgSessions[sFile]['fun'].deleteUseless(IDList, mode)
        self._freeSession(sFile)

        return 0


    async def deleteInDatabase(self, fileData: dict):
        self.fileDatabase['file'].remove(fileData)
        ret_val = await self._cleanTg(fileData['fileID'])

        if ret_val == 0:
            self.updateDatabase()


    def renameInDatabase(self, fileData: dict, newName: list):
        self.fileDatabase['file'] \
                  [self.fileDatabase['file'].index(fileData)]['rPath'] = newName
        self.fileDatabase['file'].sort(key=itemgetter('rPath'))
        self.updateDatabase()


    async def upload(self, fileData: dict, sFile: int = None):
        if not fileData['path'] or fileData['rPath'] == ['']:
            self.notification("Please enter all info")
            return
        elif not path.isfile(fileData['path']):
            self.notification("There is no file with this path")
            return

        sFile = self._useSession(sFile)
        if sFile == -1: return

        self.tgSessions[sFile]['info']['rPath'] = fileData['rPath']
        self.tgSessions[sFile]['info']['progress'] = 0
        self.tgSessions[sFile]['info']['type'] = 'upload'

        if 'index' in fileData: # resuming
            self.tgSessions[sFile]['info']['size'] = fileData['size']
        else:
            fileData['index'] = self.fileDatabase['index'][sFile] if (sFile
                                           in self.fileDatabase['index']) else 1
            fileData['chunkIndex'] = 0
            fileData['fileID'] = []
            fileData['size'] = path.getsize(fileData['path'])
            self.tgSessions[sFile]['info']['size'] = fileData['size']

            fileData['type'] = 'upload'
            sha_hash = sha256()
            hash_update = utils.asyncWrap(sha_hash.update)
            hash_digest = utils.asyncWrap(sha_hash.digest)

            with open(fileData['path'], 'rb') as f:
                for byte_block in iter(lambda: f.read(1024 * 1024), b''):
                    await hash_update(byte_block)
            fileData['checksum'] = await hash_digest()

        finalData = await self.tgSessions[sFile]['fun'].uploadFiles(fileData)

        if finalData: # finished uploading
            if len(finalData['fileData']['fileID']) > 1: # not single chunk
                self.delResumeData(sFile)

            self.fileDatabase['index'][sFile] = finalData['index']

            '''
            This is slow, a faster alternative could be bisect.insort,
            howewer, I couldn't find a way to sort by an item in dictionary
            '''
            self.fileDatabase['file'].append(finalData['fileData'])
            self.fileDatabase['file'].sort(key=itemgetter('rPath'))

            self.updateDatabase()
            self._freeSession(sFile)

        # not transferring anything
        self.tgSessions[sFile]['info']['type'] = None


    async def download(self, fileData: dict, sFile: int = None):
        if not utils.validate_path(fileData['dPath']):
            self.notification(
                "Not enough permissions for specified download path"
            )

        sFile = self._useSession(sFile) # use a free session
        if sFile == -1: return

        self.tgSessions[sFile]['info']['rPath'] = fileData['rPath']
        self.tgSessions[sFile]['info']['progress'] = 0
        self.tgSessions[sFile]['info']['size'] = fileData['size']
        self.tgSessions[sFile]['info']['type'] = 'download'

        if not 'IDindex' in fileData:
            fileData['IDindex'] = 0
            fileData['type'] = 'download'

        final_file_path = await self.tgSessions[sFile]['fun'].downloadFiles(
                                                                       fileData)

        if final_file_path: # finished downloading
            # verify checksum
            sha_hash = sha256()
            hash_update = utils.asyncWrap(sha_hash.update)
            hash_digest = utils.asyncWrap(sha_hash.digest)

            with open(final_file_path, 'rb') as f:
                for byte_block in iter(lambda: f.read(1024 * 1024), b''):
                    await hash_update(byte_block)

            act_hash = await hash_digest()
            if fileData['checksum'] != act_hash:
                await self.asyncFiles.remove(final_file_path)
                self.notif_fun("Checksum failed for {}".format(
                                                             fileData['rPath']))

            if len(fileData['fileID']) > 1: # not single-chunk
                self.delResumeData(sFile)
            self._freeSession(sFile)

        self.tgSessions[sFile]['info']['type'] = None


    async def cancelTransfer(self, sFile: int):
        if not sFile in range(1, int(self.cfg['remote']['max_sessions']) + 1):
            raise IndexError("sFile should be between 1 and {}.".format(
                int(self.cfg['remote']['max_sessions'])
            ))
        elif not self.tgSessions[sFile]['fun'].isMultiChunk():
            self.notification("Can't cancel single chunk transfers")
            return
        elif self.tgSessions[sFile]['fun'].should_stop:
            self.notification("Transfer already cancelled")
            return

        await self.tgSessions[sFile]['fun'].stop(1)


    def getUsedSessions(self) -> int:
        return int(self.cfg['remote']['max_sessions']) - len(self.freeSessions)


    def getValidResume(self) -> dict:
        valid_resume = {}

        for sFile, info in self.resumeData.items():
            if info and not self.tgSessions[sFile]['info']['type']:
                # has resume data that wasn't handled
                valid_resume[sFile] = info

        return valid_resume
